/* Copyright (c) 2025 LiuYing
*    Key Laboratory of Cyberspace Security Defense,Institute of Information Engineering, CAS
*    School of Cyber Security, University of Chinese Academy of Sciences     
*
* This file is part of the openHiTLS project.
*
* openHiTLS is licensed under the Mulan PSL v2.
* You can use this software according to the terms and conditions of the Mulan PSL v2.
* You may obtain a copy of Mulan PSL v2 at:
*
*     http://license.coscl.org.cn/MulanPSL2
*
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
* See the Mulan PSL v2 for more details.
*/
#include<stdint.h>
#include "polarlac_param.h"
#include "ntt1024_param.h"


//Testbook Multiplication for correctness test
void NormalMul_unsign_251(const uint16_t  *a, const uint16_t *s, uint16_t *b)  //常规乘法版本二
{
    int i,j;
    uint32_t m[2*DIM_N-1]={0};
    uint16_t s2[DIM_N];
    for(i=0;i<DIM_N;i++)
	{
		if(s[i]>=q_half)
			s2[i]=neg_one ; // -1 was transformed into 255 in uint8, so we recover it to 250 with mod 251
		else
			s2[i] = s[i] ;
	}
    for(i=0;i<DIM_N;i++)
    {
        for(j=0;j<DIM_N;j++)
        {
            m[i+j]+=((uint32_t)a[i]*(uint32_t)s2[j])%Q;
            m[i+j]=m[i+j]%Q;
        }
    }
  
    for(i=0;i<DIM_N-1;i++)
        b[i]=(m[i]+Q-m[i+DIM_N])%Q;
    b[DIM_N-1]=m[DIM_N-1]%Q;
    
}


#if defined LAC256
int NTT_1024(int32_t *a) 
{
    
    int t,m,i,j,s,e;
    int32_t U=0,V=0,S;
    t=DIM_N; 

    for(m=1;m<DIM_N;m<<=1) 
    {
        t=(t>>1);
        for(i=0;i<m;i++)  
        {
            s=(i*t)<<1;
            e=s+t;
            S=f_1024[m+i];
            for(j=s;j<e;j++)  
            {
                U=a[j];
                V=a[j+t]*S%NTTQ;
              
                a[j]=(U+V)%NTTQ;
                a[j+t]=(U+NTTQ-V)%NTTQ;
            }
        }
    }
    
    return 0;
}

int INTT_1024(int32_t *a) 
{
    int t,m,i,j,s,e,h;
    int32_t U=0,V=0,S;
    t=1; 
   
    for(m=DIM_N;m>1;m>>=1) 
    {
        s=0;
        h=m>>1;
        for(i=0;i<h;i++) 
        {
            e=s+t;
            S=fn_1024[h+i];
            for(j=s;j<e;j++)  
            {
                U=a[j];
                V=a[j+t];

                a[j]=(U+V)%NTTQ;

                a[j+t]=(U+NTTQ-V)*S%NTTQ;   
            }
            s=s+2*t;
        }
        t*=2;
    }
    
    for(i=0;i<DIM_N;i++)
    {
        a[i]=(a[i]*INVERSE_N)%NTTQ;
    }

    return 0;
}



int poly_mul_ntt_1024(const uint16_t  *a, const uint16_t *s, uint16_t *b)
{
    int i;
    int32_t a_buf[DIM_N], s_buf[DIM_N];
    int32_t b_buf[DIM_N];

    for(i=0;i<DIM_N;i++)
    {
        a_buf[i]=a[i];
        s_buf[i]=s[i];
    }
   
     //NTT form
    NTT_1024(a_buf);  
    NTT_1024(s_buf);  
    
     //point mul
     for(i=0;i<DIM_N;i++)
     {
         b_buf[i]=(a_buf[i]%NTTQ)*s_buf[i]%NTTQ;  
     }

     //INTT form
    INTT_1024(b_buf);

     //mod Q
     for(i=0;i<DIM_N;i++)
     {
         b[i]=b_buf[i];
     }
     
     return 0;
}

//*********************利用Montgomery Reduction求模*********************//
//32bit求模

static int32_t MontgomerymapFull(int32_t a)  //返回r=aβ^(-1) mod q, 0<r<q map到标准表示
{
    int32_t t;
    int32_t m,r;
    m = a * INVERSE_Q;
    m=m&((1<<16)-1);  // m mod β
    t = m * NTTQ;
    r = a - t;
    r >>= 16;
    return r+((r>>15)&NTTQ);
}

int NTT_Lazy_1024(int16_t *a)
{
    int t,m,i,j,s,e;
    int16_t U=0,V=0;
    uint16_t S;
    t=DIM_N;

    // m=1
    m=1;
    t=(t>>1);
        for(i=0;i<m;i++)  
        {
            s=(i*t)<<1;
            e=s+t;
            S=M_1024[m+i];
            for(j=s;j<e;j++)  
            {
                U=a[j];
                V=a[j+t];

                V=MontgomerymapFull((int32_t)V*(int32_t)S);

                a[j]=U+V-NTTQ;
                a[j+t]=U-V;
              
            }
        }

    for(m=2;m<DIM_N_2;m<<=1) 
    {
        t=(t>>1);
        for(i=0;i<m;i++)  
        {
            s=(i*t)<<1;
            e=s+t;
            S=M_1024[m+i];
            for(j=s;j<e;j++)  
            {
                U=a[j];
                V=a[j+t];

                V=MontgomerymapFull((int32_t)V*(int32_t)S);
                U = U + ((U>>15)&NTTQ);

                a[j]=(U+V-NTTQ);
                a[j+t]=(U-V);
              
            }
        }
    }

    m=DIM_N_2;
    t=(t>>1);
        for(i=0;i<m;i++)  
        {
            s=(i*t)<<1;
            e=s+t;
            S=M_1024[m+i];
            for(j=s;j<e;j++)  
            {
                U=a[j];
                V=a[j+t];

                V=MontgomerymapFull((int32_t)V*(int32_t)S);
                U = U + ((U>>15)&NTTQ);
                
                a[j]=U+V-NTTQ;
                a[j] = a[j] + ((a[j]>>15)&NTTQ);

                a[j+t]=U-V;
                a[j+t] = a[j+t] + ((a[j+t]>>15)&NTTQ);
            }
        }


    return 0;
}

int INTT_Lazy_1024(int16_t *a)
{
    int t,m,i,j,s,e,h;
    int16_t U=0,V=0;
    uint16_t S;
    t=1;
    for(m=DIM_N;m>1;m/=2)
    {
        s=0;
        h=m/2;
        for(i=0;i<h;i++)
        {
            e=s+t;
            S=Mn_1024[h+i];
            for(j=s;j<e;j++)
            {
                U=a[j];
                V=a[j+t];
                a[j] = (U+V-NTTQ);
                a[j] = a[j] + ((a[j]>>15)&NTTQ);
                a[j+t]=MontgomerymapFull((int32_t)(U-V)*(int32_t)S);

            }
            s=s+2*t;
        }
        t*=2;
    }

    for(i=0;i<DIM_N;i++)
    {
        a[i]=(a[i]*INVERSE_N_BETA)%NTTQ;
    }
    return 0;
}

int poly_mul_ntt_Lazy_1024(const uint16_t  *a, const uint16_t *s, uint16_t *b)
{
    int i;
    int16_t a_buf[DIM_N], s_buf[DIM_N],b_buf[DIM_N];

    for(i=0;i<DIM_N;i++)
    {
        a_buf[i]=a[i];
        s_buf[i]=s[i];
    }
   
     //NTT form
    NTT_Lazy_1024(a_buf);
    NTT_Lazy_1024(s_buf);
    
     //point mul
     for(i=0;i<DIM_N;i++)
     {
         b_buf[i]=MontgomerymapFull((int32_t)a_buf[i]*(int32_t)s_buf[i]);
      
     }

     //INTT form
    INTT_Lazy_1024(b_buf);

     //mod Q
     for(i=0;i<DIM_N;i++)
     {
         b[i]=(b_buf[i]+NTTQ)%NTTQ;
     }

     //没有改变传入的a和s
     return 0;
}
#endif